import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary


class SE_block(nn.Module):
    def __init__(self, in_channel, ratio):
        super(SE_block, self).__init__()
        self.squeeze_block = nn.AdaptiveAvgPool2d((1, 1))
        # can replace by linear
        self.compressed = nn.Conv2d(in_channels=in_channel, out_channels=in_channel // ratio, kernel_size=(1, 1))
        self.excitation = nn.Conv2d(in_channels=in_channel // ratio, out_channels=in_channel, kernel_size=(1, 1))
        # self.compressed = nn.Linear(in_features=in_channel, out_features=in_channel // ratio, bias=False)
        # self.excitation = nn.Linear(in_features=in_channel // ratio, out_features=in_channel, bias=False)
        self.relu = nn.ReLU(True)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        out = self.squeeze_block(x)
        out = self.compressed(out)
        out = self.relu(out)
        out = self.excitation(out)
        out = self.sigmoid(out)
        return out


def con3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=dilation, groups=groups, bias=False, dilation=dilation)


# 封装下1x1卷积层
def con1x1(in_planes, out_planes, stride=1):
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)


# resnet 50 以下使用的基本块
class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                 base_width=64, dilation=1, norm_layer=None, if_se=False):
        super(BasicBlock, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if groups != 1 or base_width != 64:
            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
        if dilation > 1:
            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")

        self.if_se = if_se
        if self.if_se:
            self.se_block = SE_block(planes, 16)

        # 下面定义BasicBlock中的各个层
        self.conv1 = con3x3(inplanes, planes, stride)
        self.bn1 = norm_layer(planes)
        # inplace为True表示进行原地操作，一般默认为False，表示新建一个变量存储操作
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = con3x3(planes, planes)
        self.bn2 = norm_layer(planes)
        self.dowansample = downsample
        self.stride = stride

    # 定义前向传播函数将前面定义的各层连接起来
    def forward(self, x):
        identity = x  # 这是由于残差块需要保留原始输入

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.if_se:
            se = self.se_block(out)
            out = out * se

        # 下采样 保证原始输入与卷积后的输出层叠加时维度相同
        if self.dowansample is not None:
            identity = self.dowansample(x)

        out = out + identity
        out = self.relu(out)

        return out


# resnet50及以上使用的基础块 可以降低参数量
class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                 base_width=64, dilation=1, norm_layer=None, if_se=False):
        """
        :param inplanes: 输入通道数
        :param planes: 输出通道数
        :param stride: 步长
        :param downsample: 下采样 适应channel大小变化
        :param groups: in_channel分组
        :param base_width: 宽度基数 默认为64
        :param dilation: 是否空洞卷积默认1不采用
        :param norm_layer: 正则化函数
        :param if_se: 是否使用SE_Block
        """
        super(Bottleneck, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = int(planes * (base_width / 64.0)) * groups

        # 定义Bottleneck中各层
        self.conv1 = con1x1(inplanes, width)
        self.norm1 = norm_layer(width)
        self.conv2 = con3x3(width, width, stride, groups, dilation)
        self.norm2 = norm_layer(width)
        self.conv3 = con1x1(width, planes * self.expansion)
        self.norm3 = norm_layer(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)

        # 适应channel大小变化 保证identity+out可行
        self.downsample = downsample
        self.stride = stride

        self.if_se = if_se
        if self.if_se:
            self.se_block = SE_block(planes * self.expansion, 16)

    def forward(self, x):
        identity = x  # 浅拷贝
        out = self.conv1(x)
        out = self.norm1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.norm2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.norm3(out)

        if self.if_se:
            se = self.se_block(out)
            out = out * se

        if self.downsample is not None:
            identity = self.downsample(x)

        out = out + identity
        out = self.relu(out)

        return out


class SE_ResNet(nn.Module):
    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
                 groups=1, width_per_group=64, replace_stride_with_dilation=None,
                 norm_layer=None, if_se=False):
        super(SE_ResNet, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self._norm_layer = norm_layer

        self.inplanes = 64
        self.dilation = 1
        if replace_stride_with_dilation is None:
            replace_stride_with_dilation = [False, False, False]
            if len(replace_stride_with_dilation) != 3:
                raise ValueError("replace_stride_with_dilation should be None "
                                 "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
        self.groups = groups
        self.base_width = width_per_group
        # 所有resnet都共有的部分
        # 将通道变为inpalnes
        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
        self.norm1 = norm_layer(self.inplanes)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

        self.layer1 = self._make_layer(block, 64, layers[0], if_se=if_se)
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
                                       dilate=replace_stride_with_dilation[0], if_se=if_se)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
                                       dilate=replace_stride_with_dilation[1], if_se=if_se)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
                                       dilate=replace_stride_with_dilation[2], if_se=if_se)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * block.expansion, num_classes)

        # 定义初始化方式
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, Bottleneck) and m.norm3.weight is not None:
                    nn.init.constant_(m.norm3.weight, 0)
                elif isinstance(m, BasicBlock) and m.norm2.weight is not None:
                    nn.init.constant_(m.norm2.weight, 0)

    def _make_layer(self, block, planes, blocks, stride=1, dilate=False, if_se=False):
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if dilate:
            self.dilation = self.dilation * stride
            stride = 1
        # 需要下采样来进行对齐
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(con1x1(self.inplanes, planes * block.expansion, stride),
                                       norm_layer(planes * block.expansion)
                                       )
        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample, self.groups, self.base_width,
                            previous_dilation, norm_layer, if_se))
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes, groups=self.groups,
                                base_width=self.base_width, dilation=self.dilation,
                                norm_layer=norm_layer, if_se=if_se))

        return nn.Sequential(*layers)

    def _forward_impl(self, x):
        x = self.conv1(x)
        x = self.norm1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.avgpool(x)

        # 接入全连接层 将输出行展开
        x = torch.flatten(x, 1)
        x = self.fc(x)

        return x

    def forward(self, x):
        return self._forward_impl(x)


def _se_resnet(arch, block, layers, pretrained, progress, **kwargs):
    # arch 网址索引
    model = SE_ResNet(block, layers, **kwargs)
    # if pretrained:
    #     return model
    return model


def se_resnet50(pretrained=False, progress=True, **kwargs):
    """
    :param kwargs:
    num_classes=1000,
    zero_init_residual=False,
    groups=1,
    width_per_group=64,
    replace_stride_with_dilation=None,
    norm_layer=None
    """
    return _se_resnet('se_resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs)


def train(model, device, epoch, train_loader, log_interval=200):
    model.train()
    running_loss = 0.0
    for i, data in enumerate(train_loader, 0):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data
        inputs = inputs.to(device)
        labels = labels.to(device)

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        # 测试模型
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        if i % log_interval == log_interval - 1:  # print every 2000 mini-batches
            print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
            running_loss = 0.0


def validate(loss_vec, acc_vec, model, test_loader, device):
    model.eval()
    val_loss, correct = 0, 0
    for data, target in test_loader:
        data = data.to(device)
        target = target.to(device)
        output = model(data)
        val_loss += criterion(output, target).data.item()
        pred = output.data.max(1)[1]
        correct += pred.eq(target.data).cpu().sum()

    val_loss /= len(test_loader)
    loss_vec.append(val_loss)

    accuracy = 100. * correct.to(torch.float32) / len(test_loader.dataset)
    acc_vec.append(accuracy)

    print('\nValidation set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        val_loss, correct, len(test_loader.dataset), accuracy))


if __name__ == "__main__":
    if torch.cuda.is_available():
        device = torch.device('cuda', 0)
    else:
        device = torch.device('cpu')

    print('Using PyTorch version:', torch.__version__,
          ' Device:', torch.cuda.get_device_name(torch.cuda.current_device()))

    # transforms.Resize((224, 224)),
    myTransforms = transforms.Compose([
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])

    transform_train = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.RandomGrayscale(),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    transform = myTransforms

    batch_size = 64  # 32 if resize
    epochs = 25  # 12 if resize
    #  num_workers origin 2 but 0 if resize
    trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                            download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                              shuffle=True, num_workers=4)

    testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                           download=True, transform=transform)
    testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                             shuffle=True, num_workers=4)

    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    net_se_resnet = se_resnet50(num_classes=10, if_se=True).to(device)
    net = net_se_resnet

    criterion = nn.CrossEntropyLoss()

    optimizer_Adam = optim.Adam(net.parameters(), lr=3e-4)
    optimizer_SGD = optim.SGD(net.parameters(), lr=0.001, momentum=0.9, weight_decay=1e-2)
    optimizer = optimizer_Adam

    loss_vec = []
    acc_vec = []

    for epoch in range(epochs):  # loop over the dataset multiple times
        train(model=net, epoch=epoch, train_loader=trainloader, device=device)
        validate(loss_vec=loss_vec, acc_vec=acc_vec, model=net, test_loader=testloader, device=device)

    print('Finished Training')

    PATH = './cifar_net.pth'
    torch.save(net.state_dict(), PATH)

    plt.figure(figsize=(5, 3))
    plt.plot(np.arange(1, epochs + 1), loss_vec)
    plt.title('validation loss')
    plt.figure(figsize=(5, 3))
    plt.plot(np.arange(1, epochs + 1), acc_vec)
    plt.title('validation accuracy')
    plt.show()
